import os
import yaml
import json

class DatasetConfigLoader:
    def __init__(self):
        self.user_cfg = ''

    # 从yml文件读取配置
    def load_yml_file(self, cfg_path):
        try:
            with open(cfg_path) as f:
                self.user_cfg = yaml.load(f, Loader=yaml.Loader)
            dataset_type = self.user_cfg.get('dataset_type')
            if dataset_type == 'VOC':
                return self._load_voc_dataset()
            elif dataset_type == 'COCO':
                return self._load_coco_dataset()
            else:
                return None
        except Exception as ex:
            print(ex)
            return None
    
    # 从cfg读取内容并转换到配置
    def load_yml_data(self, cfg):
        self.user_cfg = yaml.load(cfg, Loader=yaml.Loader)
        dataset_type = self.user_cfg.get('dataset_type')
        if dataset_type == 'VOC':
            return self._load_voc_dataset()
        elif dataset_type == 'COCO':
            return self._load_coco_dataset()
        else:
            return None

    def load_json_file(self, cfg_path):
        with open(cfg_path, 'r', encoding='utf-8') as f:
            total_cfg = json.load(f)
        self.user_cfg = total_cfg['dataset_config']
        dataset_type = self.user_cfg.get('dataset_type')
        if dataset_type == 'VOC':
            return self._load_voc_dataset()
        elif dataset_type == 'COCO':
            return self._load_coco_dataset()
        else:
            return None

    # 从json格式读取配置
    def load_json_data(self, cfg):
        self.user_cfg = cfg['dataset_config']
        dataset_type = self.user_cfg.get('dataset_type')
        if dataset_type == 'VOC':
            return self._load_voc_dataset()
        elif dataset_type == 'COCO':
            return self._load_coco_dataset()
        else:
            return None

    def _create_common_config(self, dataset_key, dataset_type):
        dataset_cfg = self.user_cfg[dataset_key]
        
        # 处理 VOC 和 COCO 不同的配置信息
        data_prefix = {'img': dataset_cfg['image_dir']} if dataset_type == 'CocoDataset' else {'sub_data_root': dataset_cfg['dataset_dir']}
        
        common_config = {
            'type': dataset_type,
            'data_root': dataset_cfg['dataset_dir'],
            'ann_file': dataset_cfg['anno_path'],
            'data_prefix': data_prefix,
            'metainfo': {'classes': self.user_cfg['classes']}
        }
        
        if dataset_type == 'XMLDataset':  # VOC 特有配置
            common_config.update({
                'img_subdir': dataset_cfg['image_dir'],
                'ann_subdir': dataset_cfg['anno_dir']
            })
        
        return {'dataset': common_config}

    def _load_voc_dataset(self):
        cfg_dict = self._base_config('XMLDataset')
        cfg_dict['val_evaluator'] = {'type': 'VOCMetric', 'metric': 'mAP', 'eval_mode': '11points'}
        cfg_dict['test_evaluator'] = {'type': 'VOCMetric', 'metric': 'mAP', 'eval_mode': '11points'}
        # print(json.dumps(cfg_dict))
        return len(self.user_cfg['classes']), cfg_dict

    def _load_coco_dataset(self):
        cfg_dict = self._base_config('CocoDataset')
        val_anno_path = os.path.join(self.user_cfg['EvalDataset']['dataset_dir'], self.user_cfg['EvalDataset']['anno_path'])
        evaluator_cfg = {'ann_file': val_anno_path, 'format_only': False, 'backend_args': None}
        cfg_dict['val_evaluator'] = evaluator_cfg
        cfg_dict['test_evaluator'] = evaluator_cfg
        # print(json.dumps(cfg_dict))
        return len(self.user_cfg['classes']), cfg_dict

    def _base_config(self, dataset_type):
        return {
            'train_dataloader': {
                'batch_size': 1,
                **self._create_common_config('TrainDataset', dataset_type)
            },
            'val_dataloader': self._create_common_config('EvalDataset', dataset_type),
            'test_dataloader': self._create_common_config('EvalDataset', dataset_type)
        }